package org.easy.bytetts;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;

public class ByteRealtimeTtsTemplate {
    private final ConcurrentHashMap<String, ByteTtsWebsocketHolder> registry = new ConcurrentHashMap<>();

    private final ByteRealtimeTtsClient byteRealtimeTtsClient;

    public ByteRealtimeTtsTemplate(ByteRealtimeTtsClient byteRealtimeTtsClient) {
        this.byteRealtimeTtsClient = byteRealtimeTtsClient;
    }

    public void init(String userId, String speaker, ByteTtsWebSocketListener listener) {
        ByteTtsWebsocketHolder byteTtsWebsocketHolder = new ByteTtsWebsocketHolder();
        byteTtsWebsocketHolder.setSpeaker(speaker);
        byteTtsWebsocketHolder.setListener(listener);
        registry.put(userId, byteTtsWebsocketHolder);
    }

    public void reInit(String userId) {
        ByteTtsWebsocketHolder holder = registry.get(userId);
        initGetSessionId(userId, holder.getSpeaker(), holder.getListener());
        holder = registry.get(userId);
        holder.setSeq(1);
    }

    private String initGetSessionId(String userId, String speaker, ByteTtsWebSocketListener listener) {
        CountDownLatch latch = new CountDownLatch(1);
        ByteTtsWebsocketHolder init = byteRealtimeTtsClient.init(speaker, listener, latch);
        try {
            latch.await();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        registry.put(userId, init);
        return userId;
    }

    public void sendMessageWithSeq(String userId, String text) {
        ByteTtsWebsocketHolder holder = registry.get(userId);
        holder.setSeq(holder.getSeq() + 1);
        byteRealtimeTtsClient.sendMessageWithSeq(userId, holder.getWebSocket(), holder.getSpeaker(), holder.getSessionId(), text, -1);
    }

    public void finishSession(String userId) {
        ByteTtsWebsocketHolder holder = registry.get(userId);

        if (holder != null && holder.getWebSocket() != null) {
            byteRealtimeTtsClient.finishSession(holder.getWebSocket(), holder.getSessionId());
        }
    }
}
